/* plantmtd.cpp */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#include <list>
#include <algorithm>

#define STRING_LEN		31
#define STRINGZ_LEN		(STRING_LEN + 1)
#define MAKESTRING2(X)	#X
#define MAKESTRING(X)	MAKESTRING2(X)
#define STRING_ALEN		MAKESTRING(STRING_LEN)

typedef unsigned int UINT;
typedef std::pair<UINT, UINT> bcont_t;
typedef char string_t[STRINGZ_LEN];
typedef struct {UINT size; string_t label;} mtdpart_t;

int main(int argc, char *argv[])
{
	const UINT first_block = 0;
	const UINT last_block = 4096; /* inclusive */
	const UINT block_size = 262144;
	const UINT div1K = 1024;
	const UINT div1M = 1048576;
	const UINT div1G = 1073741824;
	const UINT user_size = 524288000; /* 500M */

	const char mtdparts_s[] = "mtdparts=atmel_nand:256k(boot)ro,256k(original_env)ro,256k(runtime_env),256k(bootloader_env),5M(kernel1),95M(rootfs1),5M(kernel2),95M(rootfs2),100M(service),-(userdata)",
		mtdstart_s[] = "mtdparts=atmel_nand:", lsuff_s[] = "KMG";

	const int mtdstart_l = strlen(mtdstart_s);

	std::list<UINT> bblist;
	std::list<bcont_t> bcontl;
	std::list<bcont_t>::iterator bc;
	std::list<mtdpart_t> mtdparts;
	std::list<mtdpart_t>::iterator ip, it;
	UINT i, j, k, as;
	char *ep, sc;
	mtdpart_t mp;


	for (i = 1; (int)i < argc; i++) {
		UINT v = strtoul(argv[i], &ep, 10);

		if (ep == argv[i] || v < first_block || v > last_block) {
			fprintf(stderr, "usage: %s <block_id> [<block_id> ..]\n"
				"\twhere block_id is {%d..%d}\n", argv[0], first_block, last_block);
			return (EINVAL);
		}

		bblist.push_back(v);
	}

	fprintf(stdout, "----- input blocks -----\n");

	bblist.sort();
	bblist.unique();

	for (std::list<UINT>::iterator ib=bblist.begin(); ib != bblist.end(); ib++)
		fprintf(stdout, "  % 5d\n", *ib);

	fprintf(stdout, "----- continuous areas -----\n");

	for (i = first_block; i <= last_block; i = j + 1) {
		for ( ; std::binary_search(bblist.begin(), bblist.end(), i); i++);

		for (j = i + 1; j <= last_block &&
			!std::binary_search(bblist.begin(), bblist.end(), j); j++);

		bcontl.push_back(bcont_t(i, j));
	}

	for (bc = bcontl.begin(); bc != bcontl.end(); bc++)
		fprintf(stdout, "  % 5d  L%d\n", bc->first, (bc->second - bc->first));

	if (bcontl.size() && bcontl.begin()->first != 0) {
		fprintf(stderr, "block 0 is unavailable\n");
		return (EINVAL);
	}

	fprintf(stdout, "----- MTD parts -----\n");

	if (strncmp(mtdparts_s, mtdstart_s, mtdstart_l) != 0) {
		fprintf(stderr, "bad MTD string\n");
		return (EINVAL);
	}

	for (ep = (char *)mtdparts_s + mtdstart_l; ep && *ep;
		ep = ep + std::min(strcspn(ep, ",") + 1, strlen(ep))) {
		if (sscanf(ep, "%u%c(%" STRING_ALEN "s),", &mp.size, &sc, mp.label) == 3) {
			switch (sc) {
			case 'k':
			case 'K': mp.size = mp.size * div1K; break;
			case 'm':
			case 'M': mp.size = mp.size * div1M; break;
			case 'g':
			case 'G': mp.size = mp.size * div1G; break;
			default: break;
			}
			mp.label[strcspn(mp.label, ")")] = 0;
			mtdparts.push_back(mp);
		} else {
			break;
		}
	}

	if (ep && sscanf(ep, "-(%" STRING_ALEN "s),", mp.label) == 1) {
		mp.size = user_size;
		mp.label[strcspn(mp.label, ")")] = 0;
		mtdparts.push_back(mp);
	}

	if (mtdparts.empty()) {
		fprintf(stdout, "bad MTD list\n");
		return (EINVAL);
	}

	for (bc = bcontl.begin(), it = mtdparts.begin(), ip = it, j = 0;
		bc != bcontl.end() && it != mtdparts.end(); j = bc->second, bc++) {
		ip->size += block_size * (bc->first - j);
		as = (bc->second - bc->first) * block_size;

		for ( ; it != mtdparts.end() && it->size <= as; as -= it->size, ip = it++);

		ip->size += as; 
	}

	if (bc == bcontl.end() && it != mtdparts.end()) {
		fprintf(stderr, "could not place all partitions required\n");
		return (ENOSPC);
	}

	ip->size = (ip->size / div1M) * div1M;

	for (ip = mtdparts.begin(), j = 0; ip != mtdparts.end(); j += ip->size, ip++) {

		for (i = ip->size, k = 0; k < sizeof (lsuff_s) / sizeof (lsuff_s[0]); k++)
			if (i >= div1K && i % div1K == 0) i = i / div1K, sc = lsuff_s[k];

		fprintf(stdout, "  %16s  %8d%c  0x%x\n", ip->label, i, sc, j);
	}

	return (0);
}

/* end */
